import pandas as pd
import matplotlib.pyplot as plt
import matplotlib.dates as mdates

# --- 1. Load the dataset ---
# This assumes 'GSSI_model_dataset_US_v2.csv' is in the same directory.
try:
    model_df = pd.read_csv('GSSI_model_dataset_US_v2.csv', index_col=0, parse_dates=True)
except FileNotFoundError:
    print("Error: 'GSSI_model_dataset_US_v2.csv' not found.")
    # Create a dummy dataframe for the code to run without error,
    # so the user can see the plotting structure.
    dates = pd.date_range(start='2004-01-01', end='2024-12-31', freq='Q')
    model_df = pd.DataFrame(
        {'GDP_GROWTH': [0.5]*len(dates), 'MX_LX_Ratio': [1.2]*len(dates)},
        index=dates
    )
    print("Using dummy data to generate a plot template.")


# --- 2. Create the stationary GSSI series ---
# As per your analysis, the GSSI was non-stationary and you used its first difference.
model_df['Ratio_stationary'] = model_df['MX_LX_Ratio'].diff()

# Drop NaNs created by differencing for a clean plot
plot_data = model_df.dropna(subset=['GDP_GROWTH', 'Ratio_stationary'])

# --- 3. Create the Dual-Axis Plot ---
fig, ax1 = plt.subplots(figsize=(12, 7))

# Plot 1: GDP Growth (as bars, on left Y-axis)
ax1.bar(plot_data.index, plot_data['GDP_GROWTH'], 
        width=60,  # Bar width in days (approx 2 months)
        color='cornflowerblue', 
        alpha=0.7, 
        label='US GDP Growth (Q/Q %)')
ax1.set_xlabel('Date')
ax1.set_ylabel('US GDP Growth (Quarter-on-Quarter %)', color='cornflowerblue')
ax1.tick_params(axis='y', labelcolor='cornflowerblue')
ax1.axhline(0, color='gray', linestyle='--', linewidth=0.8) # Add a zero line

# Create the second Y-axis that shares the same X-axis
ax2 = ax1.twinx()

# Plot 2: GSSI (as a line, on right Y-axis)
ax2.plot(plot_data.index, plot_data['Ratio_stationary'], 
         color='crimson', 
         linewidth=2.0, 
         label='GSSI (Stationary)')
ax2.set_ylabel('GSSI (First Difference)', color='crimson')
ax2.tick_params(axis='y', labelcolor='crimson')

# --- 4. Final Touches (The "Unrejectable" Part) ---
# Title
plt.title('US GSSI vs. US GDP Growth (2004-2024)', fontsize=16)

# Format X-axis
ax1.xaxis.set_major_locator(mdates.YearLocator(4)) # Show a tick every 4 years
ax1.xaxis.set_major_formatter(mdates.DateFormatter('%Y'))

# Create unified legends
lines, labels = ax1.get_legend_handles_labels()
lines2, labels2 = ax2.get_legend_handles_labels()
ax2.legend(lines + lines2, labels + labels2, loc='upper left')

# Ensure layout is clean
plt.tight_layout()

# --- 5. Save the final file ---
plt.savefig('GSSI_vs_GDP_Plot.png')

print("Plot saved successfully as 'GSSI_vs_GDP_Plot.png'.")